import os
import fitz
import cv2
import pytesseract
from docx import Document
import tkinter as tk
from tkinter import filedialog, messagebox

# 配置 Tesseract OCR 路径，需根据实际安装情况修改
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

# 创建临时文件夹
tmp_folder = 'tmp'
if not os.path.exists(tmp_folder):
    os.makedirs(tmp_folder)


def pdf_to_images(pdf_path):
    doc = fitz.open(pdf_path)
    zoom_x = 3.0
    zoom_y = 3.0
    mat = fitz.Matrix(zoom_x, zoom_y)
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        pix = page.get_pixmap(matrix=mat)
        image_path = os.path.join(tmp_folder, f'page_{page_num}.png')
        pix.save(image_path)
    doc.close()


def preprocess_image(image_path):
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return thresh


def ocr_images():
    all_text = ""
    tessdata_dir = r"tessdata"
    # 临时设置 TESSDATA_PREFIX 环境变量
    os.environ['TESSDATA_PREFIX'] = tessdata_dir
    tessdata_dir_config = f'--tessdata-dir {tessdata_dir}'
    print(f"使用的 tessdata 路径: {tessdata_dir_config}")
    print(f"实际检查的 tessdata 文件夹路径: {tessdata_dir}")

    if not os.path.exists(tessdata_dir):
        print(f"tessdata 文件夹 {tessdata_dir} 不存在。")
        return all_text

    chi_sim_file = os.path.join(tessdata_dir, 'chi_sim.traineddata')
    eng_file = os.path.join(tessdata_dir, 'eng.traineddata')
    if not os.path.exists(chi_sim_file) or not os.path.exists(eng_file):
        print(f"语言数据文件 {chi_sim_file} 或 {eng_file} 不存在。")
        return all_text

    for filename in os.listdir(tmp_folder):
        if filename.endswith('.png'):
            image_path = os.path.join(tmp_folder, filename)
            preprocessed_img = preprocess_image(image_path)
            try:
                text = pytesseract.image_to_string(preprocessed_img, lang='chi_sim+eng', config=tessdata_dir_config)
            except Exception as e:
                print(f"OCR 识别图片 {image_path} 时出错: {e}")
                continue
            txt_filename = os.path.splitext(filename)[0] + '.txt'
            txt_path = os.path.join(tmp_folder, txt_filename)
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write(text)
            all_text += text + '\n\n'
    return all_text


def merge_txt_to_word_and_single_txt(all_text):
    doc = Document()
    for filename in sorted(os.listdir(tmp_folder)):
        if filename.endswith('.txt'):
            txt_path = os.path.join(tmp_folder, filename)
            with open(txt_path, 'r', encoding='utf-8') as f:
                text = f.read()
            doc.add_paragraph(text)
            doc.add_page_break()
    doc.save('merged_result.docx')

    with open('merged_result.txt', 'w', encoding='utf-8') as f:
        f.write(all_text)


def clean_tmp_folder():
    for filename in os.listdir(tmp_folder):
        file_path = os.path.join(tmp_folder, filename)
        try:
            if os.path.isfile(file_path):
                os.unlink(file_path)
        except Exception as e:
            print(f"Error deleting {file_path}: {e}")


def convert_pdf():
    # 从输入框中获取用户选择的 PDF 文件路径
    pdf_path = pdf_entry.get()
    # 检查路径是否为空
    if not pdf_path:
        # 如果为空，弹出错误消息框提示用户选择文件
        messagebox.showerror("错误", "请选择 PDF 文件")
        # 结束函数执行
        return
    try:
        # 调用 pdf_to_images 函数，将 PDF 文件转换为一系列 PNG 图像
        pdf_to_images(pdf_path)
        # 调用 ocr_images 函数，对生成的图像进行 OCR 识别，得到所有文本内容
        all_text = ocr_images()
        # 调用 merge_txt_to_word_and_single_txt 函数，将识别的文本合并到 Word 文档和纯文本文件中
        merge_txt_to_word_and_single_txt(all_text)
        # 调用 clean_tmp_folder 函数，清理临时文件夹中的文件
        clean_tmp_folder()
        # 弹出信息消息框，提示用户转换完成
        messagebox.showinfo("完成", "转换完成，已生成 merged_result.docx 和 merged_result.txt")
    except Exception as e:
        # 如果在转换过程中出现异常，弹出错误消息框显示具体错误信息
        messagebox.showerror("错误", f"转换过程中出现错误: {e}")


# 创建 GUI 界面
root = tk.Tk()
root.title("PDF 转文本工具")

pdf_label = tk.Label(root, text="选择 PDF 文件:")
pdf_label.pack(pady=10)

pdf_entry = tk.Entry(root, width=50)
pdf_entry.pack(pady=5)

select_button = tk.Button(root, text="选择文件", command=lambda: pdf_entry.insert(0, filedialog.askopenfilename(filetypes=[("PDF 文件", "*.pdf")])))
select_button.pack(pady=10)

convert_button = tk.Button(root, text="转换", command=convert_pdf)
convert_button.pack(pady=20)

root.mainloop()